"""
媒体流处理模块
"""

import asyncio
import json
import logging
import base64
import cv2
import numpy as np
from typing import Dict, Optional, Set, Any
import websockets
from websockets.server import WebSocketServerProtocol
from websockets.exceptions import ConnectionClosed, ConnectionClosedError
import time
import io
from av import AudioFrame, VideoFrame
import queue
import threading

logger = logging.getLogger(__name__)

class WebSocketLock:
    """WebSocket发送锁，防止并发发送问题"""
    def __init__(self):
        self._lock = asyncio.Lock()
        self._sending = False

    async def acquire(self):
        await self._lock.acquire()
        self._sending = True

    async def release(self):
        self._sending = False
        self._lock.release()

    @property
    def is_sending(self):
        return self._sending

    async def __aenter__(self):
        await self.acquire()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.release()

class MediaStreamer:
    """基于队列的线程安全音视频流处理器"""

    def __init__(self, sessionid: int):
        self.sessionid = sessionid
        self.websocket: Optional[WebSocketServerProtocol] = None
        self.is_connected = False
        self.frame_count = 0
        self.audio_frame_count = 0
        self.last_frame_time = 0
        self.last_audio_time = 0

        # 帧率控制 - 更激进的网络传输帧率优化
        self.target_fps = 15  # 进一步降低到15fps，确保稳定传输
        self.frame_interval = 1.0 / self.target_fps
        self.audio_sample_rate = 16000
        self.audio_chunk_size = 320

        # 队列系统 - 线程安全的数据传输
        self.video_queue = queue.Queue(maxsize=30)  # 增加视频队列大小
        self.audio_queue = queue.Queue(maxsize=50)  # 增加音频队列大小
        self.control_queue = queue.Queue(maxsize=50)

        # 发送任务控制
        self.sender_task = None
        self.should_stop = False

        # 性能统计
        self.start_time = time.time()
        self.frames_sent = 0
        self.audio_frames_sent = 0
        self.bytes_sent = 0

    async def register_client(self, websocket: WebSocketServerProtocol):
        """注册客户端"""
        self.websocket = websocket
        self.is_connected = True
        self.frame_count = 0
        self.audio_frame_count = 0
        self.start_time = time.time()
        self.should_stop = False

        # 启动发送任务
        self.sender_task = asyncio.create_task(self._sender_loop())

        # 发送连接确认消息
        await self.queue_control_message('connected', {
            'sessionid': self.sessionid,
            'target_fps': self.target_fps,
            'audio_sample_rate': self.audio_sample_rate
        })

    async def unregister_client(self):
        """注销客户端"""
        self.is_connected = False
        self.should_stop = True

        # 停止发送任务
        if self.sender_task:
            self.sender_task.cancel()
            try:
                await self.sender_task
            except asyncio.CancelledError:
                pass
            self.sender_task = None

        self.websocket = None

        # 清空队列
        while not self.video_queue.empty():
            try:
                self.video_queue.get_nowait()
            except:
                pass
        while not self.audio_queue.empty():
            try:
                self.audio_queue.get_nowait()
            except:
                pass

    async def _sender_loop(self):
        """在WebSocket原始事件循环中运行的发送循环"""
        try:
            iteration = 0
            while not self.should_stop and self.is_connected:
                iteration += 1
                try:
                    # 处理控制消息（高优先级）
                    while not self.control_queue.empty():
                        try:
                            message = self.control_queue.get_nowait()
                            if self.websocket:
                                if hasattr(self.websocket, 'send_json'):
                                    await self.websocket.send_json(message)
                                else:
                                    await self.websocket.send(json.dumps(message))
                                logger.debug(f"Sent control message: {message.get('control_type', 'unknown')}")
                        except queue.Empty:
                            break
                        except Exception as e:
                            logger.error(f"Control message send failed: {e}")
                            break

                    # 批量处理视频帧 - 提高发送效率
                    video_sent_this_iteration = 0
                    max_video_per_iteration = 3  # 每次最多处理3个视频帧

                    for _ in range(max_video_per_iteration):
                        try:
                            frame_data = self.video_queue.get_nowait()

                            # 预检查连接状态
                            if not self.websocket or not self.is_connected:
                                logger.info(f"WebSocket disconnected, stopping video frame sending for session {self.sessionid}")
                                return  # 直接退出sender loop

                            # 发送数据
                            try:
                                if hasattr(self.websocket, 'send_json'):
                                    await self.websocket.send_json(frame_data)
                                else:
                                    await self.websocket.send(json.dumps(frame_data))

                                self.frame_count += 1
                                self.frames_sent += 1
                                self.bytes_sent += len(frame_data.get('frame_data', ''))
                                video_sent_this_iteration += 1

                            except Exception as send_error:
                                # 检查是否是连接关闭错误
                                error_str = str(send_error).lower()
                                if any(keyword in error_str for keyword in ["close message has been sent", "connection closed", "closed"]):
                                    logger.warning(f"WebSocket connection closed detected for session {self.sessionid}")
                                    self.is_connected = False
                                    return  # 直接退出sender loop
                                else:
                                    logger.error(f"Video frame send failed: {send_error}")
                                    # 其他错误也要检查连接状态
                                    if "websocket" in error_str or "connection" in error_str:
                                        self.is_connected = False
                                        return

                        except queue.Empty:
                            break  # 队列空，跳出循环

                    # 批量处理音频帧 - 每次最多5个
                    audio_sent_this_iteration = 0
                    max_audio_per_iteration = 5

                    for _ in range(max_audio_per_iteration):
                        try:
                            audio_data = self.audio_queue.get_nowait()

                            # 预检查连接状态
                            if not self.websocket or not self.is_connected:
                                logger.info(f"WebSocket disconnected, stopping audio frame sending for session {self.sessionid}")
                                return  # 直接退出sender loop

                            # 发送数据
                            try:
                                if hasattr(self.websocket, 'send_json'):
                                    await self.websocket.send_json(audio_data)
                                else:
                                    await self.websocket.send(json.dumps(audio_data))

                                self.audio_frame_count += 1
                                self.audio_frames_sent += 1
                                self.bytes_sent += len(audio_data.get('audio_data', ''))
                                audio_sent_this_iteration += 1

                            except Exception as send_error:
                                # 检查是否是连接关闭错误
                                error_str = str(send_error).lower()
                                if any(keyword in error_str for keyword in ["close message has been sent", "connection closed", "closed"]):
                                    logger.warning(f"WebSocket connection closed detected during audio send for session {self.sessionid}")
                                    self.is_connected = False
                                    return  # 直接退出sender loop
                                else:
                                    logger.error(f"Audio frame send failed: {send_error}")
                                    # 其他错误也要检查连接状态
                                    if "websocket" in error_str or "connection" in error_str:
                                        self.is_connected = False
                                        return

                        except queue.Empty:
                            break  # 队列空，跳出循环

                    # 调试信息 - 每500次迭代打印一次状态
                    if iteration % 500 == 0:
                        queue_size_video = self.video_queue.qsize()
                        queue_size_audio = self.audio_queue.qsize()
                        logger.debug(f"Sender loop iteration {iteration} - "
                                   f"Video sent: {video_sent_this_iteration}, Audio sent: {audio_sent_this_iteration}, "
                                   f"Queue sizes - Video: {queue_size_video}, Audio: {queue_size_audio}")

                    # 动态休眠 - 如果队列为空则稍微休眠，否则不休眠
                    if self.video_queue.qsize() > 15:  # 队列压力大，不休眠全力发送
                        await asyncio.sleep(0.0001)  # 0.1ms
                    elif self.video_queue.qsize() > 5:  # 队列中等压力，短暂休眠
                        await asyncio.sleep(0.0005)  # 0.5ms
                    else:  # 队列压力小，正常休眠
                        await asyncio.sleep(0.002)   # 2ms

                except Exception as e:
                    logger.error(f"Sender loop error: {e}")
                    await asyncio.sleep(0.01)

        except asyncio.CancelledError:
            logger.info(f"Sender loop cancelled for session {self.sessionid}")
        except Exception as e:
            logger.error(f"Sender loop crashed: {e}")
        finally:
            self.is_connected = False

    async def send_video_frame(self, frame: np.ndarray) -> bool:
        """将视频帧放入发送队列（线程安全）"""
        if not self.is_connected:
            return False

        try:
            # 帧率控制
            current_time = time.time()
            if current_time - self.last_frame_time < self.frame_interval:
                return False

            # 画面优化策略
            height, width = frame.shape[:2]
            max_width = 480  # 画面宽度
            if width > max_width:
                scale = max_width / width
                new_width = max_width
                new_height = int(height * scale)
                frame = cv2.resize(frame, (new_width, new_height), interpolation=cv2.INTER_AREA)

            # 编码图像
            quality = 50  # 画面质量
            _, buffer = cv2.imencode('.jpg', frame, [
                cv2.IMWRITE_JPEG_QUALITY, quality,
                cv2.IMWRITE_JPEG_OPTIMIZE, 1
            ])

            frame_data = base64.b64encode(buffer).decode('utf-8')

            message = {
                'type': 'video_frame',
                'sessionid': self.sessionid,
                'timestamp': current_time,
                'frame_data': frame_data,
                'frame_number': self.frame_count,
                'width': frame.shape[1],
                'height': frame.shape[0],
                'quality': quality
            }

            # 智能队列管理 - 更激进的丢帧策略
            queue_size = self.video_queue.qsize()

            # 如果队列超过60%容量，跳过这一帧
            if queue_size >= int(self.video_queue.maxsize * 0.6):
                self.frame_count += 1
                return False

            # 非阻塞放入队列
            try:
                self.video_queue.put_nowait(message)
                self.last_frame_time = current_time
                self.frame_count += 1
                return True

            except queue.Full:
                # 队列满时丢弃最旧的帧
                try:
                    self.video_queue.get_nowait()
                    self.video_queue.put_nowait(message)
                    self.last_frame_time = current_time
                    return True
                except:
                    # 严重队列阻塞，完全丢弃这一帧
                    self.frame_count += 1
                    return False

        except Exception as e:
            logger.error(f"Failed to queue video frame for session {self.sessionid}: {e}")
            return False

    async def send_audio_frame(self, audio_frame: np.ndarray, eventpoint=None) -> bool:
        """将音频帧放入发送队列（线程安全）"""
        if not self.is_connected:
            return False

        try:
            # 音频帧率控制
            current_time = time.time()
            if current_time - self.last_audio_time < 0.02:  # 20ms间隔
                return False

            # 确保音频格式
            if audio_frame.dtype != np.int16:
                audio_frame = (audio_frame * 32767).astype(np.int16)

            audio_data = base64.b64encode(audio_frame.tobytes()).decode('utf-8')

            message = {
                'type': 'audio_frame',
                'sessionid': self.sessionid,
                'timestamp': current_time,
                'audio_data': audio_data,
                'sample_rate': self.audio_sample_rate,
                'channels': 1,
                'frame_size': len(audio_frame),
                'eventpoint': eventpoint
            }

            # 非阻塞放入队列
            try:
                self.audio_queue.put_nowait(message)
                self.last_audio_time = current_time
                return True
            except queue.Full:
                # 队列满时丢弃最旧的帧
                try:
                    self.audio_queue.get_nowait()
                    self.audio_queue.put_nowait(message)
                    self.last_audio_time = current_time
                    return True
                except:
                    return False

        except Exception as e:
            logger.error(f"Failed to queue audio frame for session {self.sessionid}: {e}")
            return False

    async def queue_control_message(self, control_type: str, data: Dict[str, Any] = None):
        """将控制消息放入队列"""
        message = {
            'type': 'control',
            'control_type': control_type,
            'sessionid': self.sessionid,
            'timestamp': time.time(),
            'data': data or {}
        }

        try:
            self.control_queue.put_nowait(message)
            return True
        except queue.Full:
            # 队列满时丢弃最旧的消息
            try:
                self.control_queue.get_nowait()
                self.control_queue.put_nowait(message)
                return True
            except:
                return False

    async def send_status_update(self, status: str, details: Dict[str, Any] = None):
        """发送状态更新"""
        await self.queue_control_message('status_update', {
            'status': status,
            'details': details or {}
        })

    def get_connection_stats(self) -> Dict[str, Any]:
        """获取连接统计信息"""
        duration = time.time() - self.start_time
        return {
            'sessionid': self.sessionid,
            'connected': self.is_connected,
            'duration': duration,
            'video_fps': self.frames_sent / duration if duration > 0 else 0,
            'audio_fps': self.audio_frames_sent / duration if duration > 0 else 0,
            'total_frames': self.frames_sent,
            'total_audio_frames': self.audio_frames_sent,
            'data_transferred_mb': self.bytes_sent / (1024 * 1024),
            'video_queue_size': self.video_queue.qsize(),
            'audio_queue_size': self.audio_queue.qsize()
        }

class ConnectionManager:
    """连接管理器"""

    def __init__(self):
        self.streamers: Dict[int, MediaStreamer] = {}
        self.active_connections: Set[WebSocketServerProtocol] = set()

    def create_streamer(self, sessionid: int) -> MediaStreamer:
        """创建新的流处理器"""
        streamer = MediaStreamer(sessionid)
        self.streamers[sessionid] = streamer
        return streamer

    def get_streamer(self, sessionid: int) -> Optional[MediaStreamer]:
        """获取流处理器"""
        return self.streamers.get(sessionid)

    async def remove_streamer(self, sessionid: int):
        """移除流处理器"""
        if sessionid in self.streamers:
            await self.streamers[sessionid].unregister_client()
            del self.streamers[sessionid]

    def get_all_stats(self) -> Dict[str, Any]:
        """获取所有连接的统计信息"""
        stats = {
            'total_sessions': len(self.streamers),
            'active_connections': len(self.active_connections),
            'sessions': {}
        }

        for sessionid, streamer in self.streamers.items():
            stats['sessions'][str(sessionid)] = streamer.get_connection_stats()

        return stats

    async def cleanup_disconnected(self):
        """清理已断开的连接"""
        disconnected = []
        for sessionid, streamer in self.streamers.items():
            if not streamer.is_connected:
                disconnected.append(sessionid)

        for sessionid in disconnected:
            await self.remove_streamer(sessionid)

        if disconnected:
            logger.info(f"Cleaned up {len(disconnected)} disconnected sessions")

# 全局连接管理器实例
connection_manager = ConnectionManager()